import matplotlib.pyplot as plt
import numpy as np
from matplotlib.pyplot import MultipleLocator

def func(name, m=13):
    file_path = r"power_data/{}.csv".format(name)
    arr = np.loadtxt(file_path)

    plt.rc('font', family='Times New Roman')

    n = 20
    threshold = 99.99

    x = range(1, n + 1)
    y = []
    sum = 0
    for i in x:
        sum += arr[i - 1]
        y.append(sum)

    x = np.asarray(x)
    y = np.asarray(y)
    print(y)

    # plt.xticks(np.arange(1, n + 1, 1))
    # plt.ylim(99.9, 100)

    markersize = 5

    # 选取部分范围
    start = 8
    end = 18
    x = x[start:end]
    y = y[start:end]
    sp = m-1-start

    plt.plot(x, y,
             marker='o',
             color='black',
             markersize=markersize,
             markeredgecolor='black')

    plt.plot(x[sp], y[sp],
             marker='o',
             color='red',
             markersize=markersize,
             markeredgecolor='red')

    # 划能量99.99%的线
    plt.axhline(y=threshold, ls="--", c="red")  # 添加水平直线

    # 划99.99%对应的基个数
    plt.axvline(x=m, ls="--", c="red")  # 添加垂直直线

    # 横/纵轴标题
    plt.xlabel("The number of mode", size=14)
    plt.ylabel("Percentage of modal energy", size=14)

    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.get_xticklabels()[sp+1].set_color("red")  # 只设置特定刻度的颜色
    ax.spines['bottom'].set_linewidth(2)  # 设置底部坐标轴的粗细
    ax.spines['left'].set_linewidth(2)

    # 刻度
    ax = plt.gca()

    y_major_locator = MultipleLocator(0.01)  # 把y轴的刻度间隔设置为0.01，并存在变量里
    # 把y轴的主刻度设置为0.01的倍数
    ax.yaxis.set_major_locator(y_major_locator)

    x_major_locator = MultipleLocator(1)  # 把y轴的刻度间隔设置为0.01，并存在变量里
    # 把y轴的主刻度设置为0.01的倍数
    ax.xaxis.set_major_locator(x_major_locator)

    # plt.show()

    plt.savefig('power_data/{}.png'.format(name), dpi=1000, format='png', bbox_inches='tight')
    plt.close()

if __name__ == "__main__":
    func(name="POD_eigenvalue_percentage", m=13)
    # func(name="POD_eigenvalue_percentage_GBDT", m=13)
    # func(name="POD_eigenvalue_percentage_LGB", m=12)
    # func(name="POD_eigenvalue_percentage_NOML", m=13)
    # func(name="POD_eigenvalue_percentage_XGB", m=12)
